import numpy as np

import sko
from sko.operators import mutation, crossover
# 进化算法
from sko.operators_gpu import mutation_gpu, crossover_gpu, selection_gpu


class GA_MTSP(sko.GA.GA_TSP):
    mutation = mutation.mutation_vrptw1
    crossover = crossover.muti_crossover_pmx2

    def __init__(self, func, args):
        self.args = args
        self.has_constraint = False
        self.endpoint = args.endpoint
        super().__init__(func, args.n_dim, size_pop=args.batch, max_iter=args.epochs, prob_mut=args.mut_std,
                         early_stop=args.max_stay_step)

        self.gpu_selection = selection_gpu.selection_tournament_faster
        self.gpu_crossover = crossover_gpu.crossover_pmx_vrptw
        self.gpu_mutation = mutation_gpu.mutation_vrptw

    def crtbp(self):
        # create the population
        tmp = np.random.rand(self.size_pop, self.len_chrom)
        self.Chrom = np.zeros([self.size_pop, self.endpoint, self.endpoint + self.n_dim])
        self.Chrom[:, np.arange(self.args.n_dim) % self.endpoint,
        np.arange(self.args.n_dim) // self.endpoint] = tmp.argsort(axis=1) + 1
        self.Chrom = self.Chrom.astype(np.int16)

        assert np.all(np.sum(self.Chrom != 0, (-1, -2)) == self.args.n_dim)

        tmpChrom_index = np.roll(self.Chrom, shift=1, axis=-1) != 0
        tmpChrom_index[:, :, 0] = True
        tmpChrom_after = self.Chrom[tmpChrom_index].reshape(self.args.batch, -1)
        self.Chrom = self.Chrom = np.roll(tmpChrom_after, shift=1, axis=1)

        return self.Chrom

    def TopN(self, n):
        v, indices1 = np.unique(self.generation_best_Y, return_index=True)
        indices2 = np.argsort(v)
        n = min(indices1.shape[0], n)
        indices = indices1[indices2[:n]]

        return np.array(self.generation_best_X)[indices], np.array(self.generation_best_Y)[indices]
